from typing import Dict, List
import numpy as np
import pandas as pd

class SymbolPool:
    """标的池管理类"""
    
    def __init__(self, capacity: int, correlation_window: int):
        self.capacity = capacity
        self.correlation_window = correlation_window
        self.symbols: Dict[str, pd.DataFrame] = {}
        
    def add_symbol(self, symbol: str, data: pd.DataFrame) -> None:
        """添加标的到池中"""
        if len(self.symbols) >= self.capacity:
            raise ValueError(f"标的池已满，最大容量为{self.capacity}")
        self.symbols[symbol] = data
        
    def remove_symbol(self, symbol: str) -> None:
        """从池中移除标的"""
        if symbol in self.symbols:
            del self.symbols[symbol]
            
    def calculate_correlations(self) -> pd.DataFrame:
        """计算标的间相关性矩阵"""
        if not self.symbols:
            return pd.DataFrame()
            
        returns = pd.DataFrame({
            sym: data['returns'].iloc[-self.correlation_window:]
            for sym, data in self.symbols.items()
        })
        return returns.corr()